import numpy as np
import pandas as pd
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt

# 读取数据
data = pd.read_excel('./assets/data.xlsx')

# 获取有效数据
years = data.columns[1:].astype(int)
visitor = pd.to_numeric(data.iloc[0, 1:], errors='coerce').values
investment = pd.to_numeric(data.iloc[4, 1:], errors='coerce').values
satisfaction = pd.to_numeric(data.iloc[1, 1:], errors='coerce').values

# 对 satisfaction 数据进行插值，参考游客数量
interp_years = np.arange(years.min(), years.max() + 1, 1)
interp_satisfaction = np.interp(interp_years, years[~np.isnan(satisfaction)], satisfaction[~np.isnan(satisfaction)])

# 删除 NaN 值
valid_indices = ~np.isnan(visitor) & ~np.isnan(investment)
valid_years = years[valid_indices]
valid_visitors = visitor[valid_indices]
valid_investment = investment[valid_indices]

# 数据归一化
visitor_mean = np.mean(valid_visitors)
visitor_std = np.std(valid_visitors)
investment_mean = np.mean(valid_investment)
investment_std = np.std(valid_investment)

normalized_visitors = (valid_visitors - visitor_mean) / visitor_std
normalized_investment = (valid_investment - investment_mean) / investment_std

# 打印插值后的满意度数据
print("Interpolated Satisfaction Data:")
print(interp_satisfaction)

# 定义满意度模型函数
def satisfaction_model(t, a, b, c):
    N_t = np.interp(t, valid_years, normalized_visitors)  # 插值后的游客数量
    C_t = np.interp(t, valid_years, normalized_investment)  # 插值后的投资数据
    return a * N_t + b * C_t + c

# 提供初始参数估计值
initial_params = [0.1, 0.1, 0.1]

# 拟合数据
popt, pcov = curve_fit(satisfaction_model, interp_years, interp_satisfaction, p0=initial_params)
a, b, c = popt

# 打印拟合参数
print(f"a: {a:.6f}")
print(f"b: {b:.6f}")
print(f"c: {c:.6f}")

# 绘制拟合结果
plt.scatter(interp_years, interp_satisfaction, color='blue', label='Interpolated Data')
plt.plot(interp_years, satisfaction_model(interp_years, *popt), color='red', label='Fitted Curve')
plt.xlabel('Years')
plt.ylabel('Satisfaction')
plt.title('Satisfaction Model Fit')
plt.legend()
plt.show()